import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


# ===== 数据获取模块 =====
def get_etf_list():
    """获取A股所有ETF列表（网页8方法）"""
    df = ak.stock_zh_a_spot_em()
    etf_list = df[df['名称'].str.contains('ETF')]['代码'].tolist()
    return etf_list[:10]  # 取前10只ETF测试


def download_etf_data(symbol):
    """获取ETF历史数据（网页8方法）"""
    df = ak.stock_zh_a_hist(symbol=symbol, period="daily", adjust="hfq")
    # 检查列名并进行相应的重命名
    if '日期' in df.columns:
        df.rename(columns={
            '日期': 'date', '开盘': 'open', '收盘': 'close',
            '最高': 'high', '最低': 'low', '成交量': 'volume'
        }, inplace=True)
    elif 'date' in df.columns:
        df.rename(columns={
            '开盘': 'open', '收盘': 'close',
            '最高': 'high', '最低': 'low', '成交量': 'volume'
        }, inplace=True)
    else:
        raise ValueError("Unexpected column names in the data fetched from akshare")

    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    return df


# ===== 策略核心模块 =====
class MultiFactorStrategy:
    def __init__(self, df):
        self.df = df
        self.calculate_indicators()

    def calculate_indicators(self):
        """计算三大指标（网页9/12/16逻辑）"""
        # RSI指标
        self.df['daily_rsi'] = self.calculate_rsi(self.df['close'], 14)
        weekly_df = self.df.resample('W-FRI').last()
        weekly_df['weekly_rsi'] = self.calculate_rsi(weekly_df['close'], 14)
        self.df = self.df.join(weekly_df['weekly_rsi'], how='left').ffill()

        # 布林带指标
        self.df['middle'] = self.df['close'].rolling(20).mean()
        self.df['std'] = self.df['close'].rolling(20).std()
        self.df['upper'] = self.df['middle'] + 2 * self.df['std']
        self.df['lower'] = self.df['middle'] - 2 * self.df['std']
        self.df['std_pct_change'] = self.df['std'].pct_change(5)  # 5日标准差变化

        # OBV指标
        self.df['obv'] = self.calculate_obv(self.df['close'], self.df['volume'])
        self.df['obv_ma30'] = self.df['obv'].rolling(30).mean()
        self.df['obv_ema5'] = self.df['obv'].ewm(span=5).mean()
        self.df['obv_ema30'] = self.df['obv'].ewm(span=30).mean()

    def calculate_rsi(self, series, window=14):
        delta = series.diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
        rs = gain / loss
        return 100 - (100 / (1 + rs))

    def calculate_obv(self, close, volume):
        obv = (np.sign(close.diff()) * volume).fillna(0).cumsum()
        return obv

    def generate_signals(self):
        """生成交易信号（网页9/11逻辑）"""
        # 条件判断
        cond_rsi = (self.df['daily_rsi'] > 50) & (self.df['weekly_rsi'] > 60)
        cond_boll = (self.df['close'] > self.df['middle']) & (self.df['std_pct_change'] > 0.15)
        cond_obv = (self.df['obv'] > self.df['obv_ma30']) & (self.df['obv_ema5'] > self.df['obv_ema30'])

        # 信号生成（至少满足两个条件）
        self.df['signal'] = 0
        self.df.loc[cond_rsi & cond_boll, 'signal'] = 1
        self.df.loc[cond_rsi & cond_obv, 'signal'] = 1
        self.df.loc[cond_boll & cond_obv, 'signal'] = 1
        return self.df


# ===== 回测引擎模块 =====
class Backtester:
    def __init__(self, capital=100000):
        self.capital = capital
        self.position = 0
        self.trades = []

    def run(self, df):
        df['returns'] = df['close'].pct_change()
        df['strategy'] = 1  # 初始化策略净值

        for i in range(1, len(df)):
            # 开仓逻辑（网页11风控）
            if df.iloc[i - 1]['signal'] == 1 and self.position == 0:
                # 量价筛选（网页11标准）
                if df.iloc[i - 20:i]['volume'].mean() < 1e8:  # 1亿成交额过滤
                    continue

                self.position = 0.3  # 30%仓位
                entry_price = df.iloc[i]['open']
                self.trades.append({'date': df.index[i], 'action': 'buy', 'price': entry_price})

            # 平仓逻辑
            elif self.position > 0:
                # 止损逻辑（网页9标准）
                stop_loss_price = df.iloc[i - 1]['lower'] * 0.95
                if df.iloc[i]['low'] < stop_loss_price:
                    exit_price = stop_loss_price
                    self.position = 0
                    self.trades.append({'date': df.index[i], 'action': 'sell', 'price': exit_price})

                # 止盈逻辑
                elif df.iloc[i]['close'] > df.iloc[i - 1]['upper']:
                    exit_price = df.iloc[i]['close']
                    self.position = 0
                    self.trades.append({'date': df.index[i], 'action': 'sell', 'price': exit_price})

            # 更新净值
            df.iloc[i, df.columns.get_loc('strategy')] = df.iloc[i - 1]['strategy'] * (
                        1 + self.position * df.iloc[i]['returns'])

        return df, pd.DataFrame(self.trades)


# ===== 执行回测 =====
if __name__ == "__main__":
    # 获取ETF数据（网页5示例）
    symbol = '563330'  # 华泰柏瑞中证A股ETF[5](@ref)
    df = download_etf_data(symbol)

    # 执行策略
    strategy = MultiFactorStrategy(df)
    df = strategy.generate_signals()

    # 回测运行
    backtester = Backtester(capital=100000)
    result_df, trades_df = backtester.run(df)

    # 绩效分析
    print(f"最终净值: {result_df['strategy'].iloc[-1]:.2f}")
    print("交易记录:")
    print(trades_df)

    # 可视化
    plt.figure(figsize=(12, 6))
    result_df[['close', 'strategy']].plot(secondary_y='strategy')
    plt.title('Multi-Factor Strategy Performance')
    plt.show()